# -*- coding: utf-8 -*-
import cv2
import json
from predictor_wrapper import PaddlePaddlePredictor
import numpy as np
import os
import platform

MISSION_NUM = 11

class Detector:
    """ base class for Detector interface"""

    def __init__(self, model_dir, input_shape=[1, 3, 480, 480]):
        if platform.system() == "Windows":  #PC平台，用于测试与验证
            self.predictor = PaddlePaddlePredictor()
        elif platform.release().startswith('4.14.0'):  #EdgeBoard FZ3B，完全模型竞速组ROM，1.8.1软核
            from predictor_wrapper import EdgeBoard181Predictor
            self.predictor = EdgeBoard181Predictor()
        else: #EdgeBoard FZ3A，创意组ROM，1.5.1软核
            from predictor_wrapper import PaddleLitePredictor
            self.predictor = PaddleLitePredictor()

        self.threshold, self.label_list = self.load_model(model_dir)
        self.input_shape = input_shape

    def load_model(self, model_dir):
        self.predictor.load(model_dir)
        json_path = model_dir + "/model.json"
        with open(json_path) as j_file:
            j_data = json.load(j_file)
        return j_data["threshold"], j_data["label_list"]

    def image_normalization(self, image):
        shape = self.input_shape
        img = cv2.resize(image, (shape[3], shape[2]))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        img = img.astype(np.float32)
        img -= 127.5
        img *= 0.007843
        if platform.system() == "Windows":
            img = img.transpose((2, 0, 1))  # HWC to CHW
        img = img.reshape(1, 3, shape[3], shape[2])
        return img

    def infer_ssd(self, data):
        self.predictor.set_input(data, 0)
        self.predictor.run()
        out = self.predictor.get_output(0)
        return np.array(out)

    def res_to_detection(self, item, shape):
        detection_object = DetectionResult()
        detection_object.index = int(item[0])
        detection_object.score = item[1]
        detection_object.name = self.label_list[int(item[0])]
        box = item[2:6] * np.array([shape[1], shape[0], shape[1], shape[0]])

        box[0] = max(box[0], 0)
        box[1] = max(box[1], 1)
        box[2] = min(box[2], shape[1])
        box[3] = min(box[3], shape[0])
        # box[0] = 0 if box[0] <= 0 else box[0]
        # box[1] = 0 if box[1] <= 0 else box[1]
        # box[2] = shape[1] if box[2] >= shape[1] else box[2]
        # box[3] = shape[0] if box[3] >= shape[0] else box[3]

        detection_object.box = box.astype(dtype=int)
        return detection_object

    def detect(self, image):
        data = self.image_normalization(image)
        out = self.infer_ssd(data)
        # print(out)
        results = []
        if len(out) <= 0:
            return results
        # out = np.array(out)
        # 设置numpy打印的格式
        np.set_printoptions(precision=4, suppress=True)
        predict_test = np.array([data for data in out if data[1] > self.threshold])
        # print("predict_test=", predict_test)

        if len(predict_test) <= 0 or predict_test[0][0] < 0:
            return results

        count = 0
        max_indexes = [-1] * MISSION_NUM
        max_scores = [-1] * MISSION_NUM
        for label, score in predict_test[:, 0:2]:  # zip(predict_data, predict_score):
            # print(label, score)
            if score > max_scores[int(label)]:
                max_indexes[int(label)] = count
                max_scores[int(label)] = score
            count += 1
        selected_indexes = [i for i in max_indexes if i != -1]
        task_index = selected_indexes
        res = predict_test[task_index, :]
        for item in res:
            results.append(self.res_to_detection(item, image.shape))

        return results



class DetectionResult:
    def __init__(self, index=0, score=0, name="", box=None, shape=None):
        if shape is None:
            shape = [480, 640]
        if box is None:
            box = [0, 0, 0, 0]
        self.index = index
        self.score = score
        self.name = name
        self.box = list(map(int, box))
        self.center_x = int((self.box[0] + self.box[2]) / 2)
        self.center_y = int((self.box[1] + self.box[3]) / 2)
        self.shape = shape

    def __repr__(self):
        self.center_x = int((self.box[0] + self.box[2]) / 2)
        self.center_y = int((self.box[1] + self.box[3]) / 2)
        width = self.box[2] - self.box[0]
        height = self.box[3] - self.box[1]
        return "index:{} name:{} score:{:.4f} center:({},{}), width:{} height:{}".format(self.index, self.name,
                                                                                         self.score, self.center_x,
                                                                                         self.center_y, width, height)